# 导包
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_curve, roc_auc_score
# 导入XGBoost
from xgboost import XGBClassifier
data = pd.read_csv('../data/train.csv')
data.info()
# 删除无用列
data = data.drop(['StandardHours'], axis=1)

#  数据预处理
x = data.drop(['Attrition'], axis=1)
y = data['Attrition']
x = pd.get_dummies(x)
x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=42)
# 训练模型
model = XGBClassifier(n_estimators=100)
model.fit(x_train, y_train)
# 模型评估
print(model.score(x_test, y_test))
print('-----------------------------')
# 预测
y_pred = model.predict(x_test)
print(roc_auc_score(y_test, y_pred))
print('-----------------------------')

